import os
import glob
import argparse
from pathlib import Path
from gmssl import *
import gmalg

def verify_keys (prikey, pubkey):
    sm2 = gmalg.SM2()
    pk  = sm2.generate_pk(prikey)
    assert pubkey == pk, "公钥验证失败"

def sm2_dec(prikey_file, pubkey_file, plain_file, cipher_file):
    prikey = read_binary_file_(prikey_file)
    pubkey = read_binary_file_(pubkey_file)
    plain = read_binary_file_(plain_file)
    cipher = read_binary_file_(cipher_file)

    verify_keys (prikey, pubkey)

    sm2 = gmalg.SM2(prikey, pubkey)
    assert sm2.decrypt(cipher) == plain, "验证失败"
    #print("ciphertext", cipher.hex())
    #print("plaintext", plain.hex())

def read_binary_file_(file_path):
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        return binary_data

def read_binary_file(file_path):
    """读取二进制文件内容并转换为字符串"""
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        # 尝试解码为字符串，如果失败则用十六进制表示
        return binary_data.hex().upper()

def main():
    parser = argparse.ArgumentParser(description='遍历目录中的二进制文件')
    parser.add_argument('directory', help='目标目录路径')
    parser.add_argument('mode', choices=['enc', 'dec'], default='enc', help='运行模式')
    args = parser.parse_args()
    directory = args.directory
    mode = args.mode

    if not os.path.isdir(directory):
        print(f"错误：目录不存在 - {directory}")
        return
      
    # 创建输出内容
    merged_content = ""
    
    # 遍历0-9，合并对应的公钥和私钥文件
    for i in range(10):
        if mode == 'enc' :
            pubkey_file = os.path.join(directory, f'sm2_pubkey_{i}')
            prikey_file = os.path.join(directory, f'sm2_prikey_{i}')
            plain_file = os.path.join(directory, f'sm2_plain_{i}')
            cipher_file = os.path.join(directory, f'sm2_cipher_{i}')
            
            sm2_dec(prikey_file, pubkey_file, plain_file, cipher_file)
            # 读取文件内容
            pubkey_content = read_binary_file(pubkey_file)
            prikey_content = read_binary_file(prikey_file)
            plain_content = read_binary_file(plain_file)
            cipher_content = read_binary_file(cipher_file)
            length = int(len(plain_content) / 2) 
            #print(f"plain len {length}")
            
            # 添加到合并内容
            merged_content += f"公钥= {pubkey_content}\n"
            merged_content += f"私钥= {prikey_content}\n"
            merged_content += f"明文长度= {length:08x}\n"
            merged_content += f"明文= {plain_content}\n"
            merged_content += f"密文= {cipher_content}\n"
            merged_content += "\n"
        else:
            pubkey_file = os.path.join(directory, f'sm2_pubkey_{i}')
            prikey_file = os.path.join(directory, f'sm2_prikey_{i}')
            plain_file = os.path.join(directory, f'sm2_replain_{i}')
            cipher_file = os.path.join(directory, f'sm2_cipher_{i}')
            
            sm2_dec(prikey_file, pubkey_file, plain_file, cipher_file)
            # 读取文件内容
            pubkey_content = read_binary_file(pubkey_file)
            prikey_content = read_binary_file(prikey_file)
            plain_content = read_binary_file(plain_file)
            cipher_content = read_binary_file(cipher_file)
            length = int(len(plain_content) / 2) 
            #print(f"plain len {length}")
            
            # 添加到合并内容
            merged_content += f"公钥= {pubkey_content}\n"
            merged_content += f"私钥= {prikey_content}\n"
            merged_content += f"明文长度= {length:08x}\n"
            merged_content += f"密文= {cipher_content}\n"
            merged_content += f"明文= {plain_content}\n"
            merged_content += "\n"

    
    # 写入到txt文件
    if mode == 'enc':
        output_file = 'sm2加密.txt'
    else:
        output_file = 'sm2解密.txt'

    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(merged_content)
    
    print(f"合并完成，结果已保存到 {output_file}")

if __name__ == "__main__":
    main()
